package insurance.udf

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._

object udaf_test extends UserDefinedAggregateFunction {
    // 聚合函数的输入数据结构
    override def inputSchema: StructType = new StructType()
            .add(StructField("qx", DoubleType))
            .add(StructField("lx", DoubleType))
            .add(StructField("bpp_", IntegerType))

    // 缓存区数据结构
    override def bufferSchema: StructType = new StructType()
            .add(StructField("old_qx", DoubleType))
            .add(StructField("old_lx", DoubleType))

    // 聚合函数返回值数据结构
    override def dataType: DataType = DoubleType

    // 聚合函数是否是幂等的，即相同输入是否总是能得到相同输出
    override def deterministic: Boolean = true

    // 初始化缓冲区
    override def initialize(buffer: MutableAggregationBuffer) = {
        buffer(0) = null
        buffer(1) = null
    }

    // 分组后的数据，会一行一行地依次进入这个update方法进行处理
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {

        if (null == input.get(1)) {
            //上一条 old_qx
            val old_qx: Double = buffer.getDouble(0)
            //上一条 old_lx
            val old_lx: Double = buffer.getDouble(1)
            // 当前 qx
            val this_qx: Double = input.getDouble(0)
            //计算 当前条 lx
            val this_lx: Double = old_lx * (1 - old_qx) * input.getInt(2)

            //更新 缓存区
            buffer(0) = this_qx
            buffer(1) = this_lx
        } else {
            //只有当policy_year=1时，才会进来。
            //将第一轮次的qx,lx直接分别更新到缓存中。
            buffer(0) = input.get(0)
            buffer(1) = input.get(1)
        }

    }

    // 如果有多个分区，那么每个分区都会维护一个缓冲区，分区间如何合并缓冲区成为一个缓冲区
// 合并聚合函数缓冲区，如果用作窗口函数且有order by排序的话，只会有一个分区，不用考虑合并，方法体直接给个null。
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = null

    override def evaluate(buffer: Row): Any = {
        buffer.getDouble(1)
    }
}
